import { UserOutlined } from '@ant-design/icons';
import { Welcome, Bubble, Sender, useXAgent, useXChat, XRequest } from '@ant-design/x';
import { Flex } from 'antd';
import React from 'react';

const BASE_URL = 'https://api.x.ai';
const PATH = '/v1/chat/completions';
const MODEL = 'grok-beta';
const API_KEY = 'Bearer xAI的KEY';

const exampleRequest = XRequest({
  baseURL: BASE_URL + PATH,
  model: MODEL,
  dangerouslyApiKey: API_KEY
});

const roles = {
  ai: {
    placement: 'start' as const, // 显式指定为 'start'
    avatar: {
      icon: <UserOutlined />,
      style: {
        background: '#fde3cf',
      },
    },
    typing: {
      step: 5,
      interval: 20,
    },
    style: {
      maxWidth: 600,
    },
  },
  local: {
    placement: 'end' as const, // 显式指定为 'end'
    avatar: {
      icon: <UserOutlined />,
      style: {
        background: '#87d068',
      },
    },
  },
};
let mockSuccess = false;

interface Message {
  choices: Array<{
    message: {
      content: string;
    };
  }>;
}

const App = () => {
  const [content, setContent] = React.useState('');
  const [status, setStatus] = React.useState<string | undefined>();
  const [lines, setLines] = React.useState<Message[]>([]);

  const [agent] = useXAgent({
    request: async ({ message }, { onSuccess, onError }) => {
      setStatus('pending');
      await exampleRequest.create(
        {
          messages: [
            {
              role: 'user',
              content: message,
            },
          ],
          stream: false,
          temperature: 0
        },
        {
          onSuccess: (messages) => {
            setStatus('success');
            console.log('onSuccess', messages);
          },
          onError: (error) => {
            setStatus('error');
            console.error('onError', error);
          },
          onUpdate: (msg: Message) => {
            setLines((pre) => [...pre, msg]);
            console.log('onUpdate', msg);
            onSuccess(`xAI: ${msg.choices[0].message.content}`); // xAI成功返回
          },
        },
      );

      onError(new Error('Mock request failed')); // 模拟失败返回
    },
  });

  const { onRequest, messages } = useXChat({
    agent,
    requestPlaceholder: 'Waiting...', // 请求占位符
    requestFallback: 'Mock failed return. Please try again later.', // 请求失败时的回退消息
  });

  return (
    <Flex vertical gap="middle"
      style={{
        padding: '24px',
      }}
    >
      <Welcome
        icon="https://mdn.alipayobjects.com/huamei_iwk9zp/afts/img/A*s5sNRo5LjfQAAAAAAAAAAAAADgCCAQ/fmt.webp"
        title="Hello, I'm AntX-O"
        description="Base on Ant Design X & Ollama3.1"
      />
      <Bubble.List
        roles={roles}
        style={{
          maxHeight: '68vh',
        }}
        items={messages.map(({ id, message, status }) => ({
          key: id,
          loading: status === 'loading',
          role: status === 'local' ? 'local' : 'ai',
          content: message,
        }))}
      />
      <Sender
        loading={agent.isRequesting()}
        value={content}
        onChange={setContent}
        onSubmit={(nextContent) => {
          onRequest(nextContent);
          setContent('');
        }}
      />
    </Flex>
  );
};

export default App;